// Copyright Envoy AI Gateway Authors
// SPDX-License-Identifier: Apache-2.0
// The full text of the Apache license is available in the LICENSE file at
// the root of the repo.

package translator

import (
	"strconv"
	"strings"
	"testing"

	"github.com/stretchr/testify/require"

	anthropicschema "github.com/envoyproxy/ai-gateway/internal/apischema/anthropic"
	"github.com/envoyproxy/ai-gateway/internal/internalapi"
)

func TestAnthropicToAnthropic_RequestBody(t *testing.T) {
	for _, tc := range []struct {
		name              string
		original          []byte
		body              anthropicschema.MessagesRequest
		forceBodyMutation bool
		modelNameOverride string

		expRequestModel internalapi.RequestModel
		expNewBody      []byte
	}{
		{
			name:              "no mutation",
			original:          []byte(`{"model":"claude-2","messages":[{"role":"user","content":"Hello!"}]}`),
			body:              anthropicschema.MessagesRequest{Stream: false, Model: "claude-2"},
			forceBodyMutation: false,
			modelNameOverride: "",
			expRequestModel:   "claude-2",
			expNewBody:        nil,
		},
		{
			name:              "model override",
			original:          []byte(`{"model":"claude-2","messages":[{"role":"user","content":"Hello!"}], Stream: true}`),
			body:              anthropicschema.MessagesRequest{Stream: true, Model: "claude-2"},
			forceBodyMutation: false,
			modelNameOverride: "claude-100.1",
			expRequestModel:   "claude-100.1",
			expNewBody:        []byte(`{"model":"claude-100.1","messages":[{"role":"user","content":"Hello!"}], Stream: true}`),
		},
		{
			name:              "force mutation",
			original:          []byte(`{"model":"claude-2","messages":[{"role":"user","content":"Hello!"}]}`),
			body:              anthropicschema.MessagesRequest{Stream: false, Model: "claude-2"},
			forceBodyMutation: true,
			modelNameOverride: "",
			expRequestModel:   "claude-2",
			expNewBody:        []byte(`{"model":"claude-2","messages":[{"role":"user","content":"Hello!"}]}`),
		},
	} {
		t.Run(tc.name, func(t *testing.T) {
			translator := NewAnthropicToAnthropicTranslator("", tc.modelNameOverride)
			require.NotNil(t, translator)

			headerMutation, bodyMutation, err := translator.RequestBody(tc.original, &tc.body, tc.forceBodyMutation)
			require.NoError(t, err)
			expHeaders := []internalapi.Header{
				{pathHeaderName, "/v1/messages"},
			}
			if bodyMutation != nil {
				expHeaders = append(expHeaders, internalapi.Header{contentLengthHeaderName, strconv.Itoa(len(bodyMutation))})
			}
			require.Equal(t, expHeaders, headerMutation)
			require.Equal(t, tc.expNewBody, bodyMutation)

			require.Equal(t, tc.expRequestModel, translator.(*anthropicToAnthropicTranslator).requestModel)
			require.Equal(t, tc.body.Stream, translator.(*anthropicToAnthropicTranslator).stream)
		})
	}
}

func TestAnthropicToAnthropic_ResponseHeaders(t *testing.T) {
	translator := NewAnthropicToAnthropicTranslator("", "")
	require.NotNil(t, translator)

	headerMutation, err := translator.ResponseHeaders(nil)
	require.NoError(t, err)
	require.Nil(t, headerMutation)
}

func TestAnthropicToAnthropic_ResponseBody_non_streaming(t *testing.T) {
	translator := NewAnthropicToAnthropicTranslator("", "")
	require.NotNil(t, translator)
	const responseBody = `{"model":"claude-sonnet-4-5-20250929","id":"msg_01J5gW6Sffiem6avXSAooZZw","type":"message","role":"assistant","content":[{"type":"text","text":"Hi! 👋 How can I help you today?"}],"stop_reason":"end_turn","stop_sequence":null,"usage":{"input_tokens":9,"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"cache_creation":{"ephemeral_5m_input_tokens":0,"ephemeral_1h_input_tokens":0},"output_tokens":16,"service_tier":"standard"}}`

	headerMutation, bodyMutation, tokenUsage, responseModel, err := translator.ResponseBody(nil, strings.NewReader(responseBody), true, nil)
	require.NoError(t, err)
	require.Nil(t, headerMutation)
	require.Nil(t, bodyMutation)
	expected := tokenUsageFrom(9, 0, 16, 25)
	require.Equal(t, expected, tokenUsage)
	require.Equal(t, "claude-sonnet-4-5-20250929", responseModel)
}

func TestAnthropicToAnthropic_ResponseBody_streaming(t *testing.T) {
	translator := NewAnthropicToAnthropicTranslator("", "")
	require.NotNil(t, translator)
	translator.(*anthropicToAnthropicTranslator).stream = true

	// We split the response into two parts to simulate streaming where each part can end in the
	// middle of an event.
	const responseHead = `event: message_start
data: {"type":"message_start","message":{"model":"claude-sonnet-4-5-20250929","id":"msg_01BfvfMsg2gBzwsk6PZRLtDg","type":"message","role":"assistant","content":[],"stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":9,"cache_creation_input_tokens":0,"cache_read_input_tokens":1,"cache_creation":{"ephemeral_5m_input_tokens":0,"ephemeral_1h_input_tokens":0},"output_tokens":0,"service_tier":"standard"}}    }

event: content_block_start
data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}      }

event: content_block_delta
data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hi"}           }

event: content_block_delta
data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"! 👋 How"}      }

`

	const responseTail = `
event: ping
data: {"type": "ping"}

event: content_block_delta
data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":" can I help you today?"}   }

event: content_block_stop
data: {"type":"content_block_stop","index":0             }

event: message_delta
data: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":16}               }

event: message_stop
data: {"type":"message_stop"       }`

	headerMutation, bodyMutation, tokenUsage, responseModel, err := translator.ResponseBody(nil, strings.NewReader(responseHead), false, nil)
	require.NoError(t, err)
	require.Nil(t, headerMutation)
	require.Nil(t, bodyMutation)
	expected := tokenUsageFrom(10, 1, 0, 10)
	require.Equal(t, expected, tokenUsage)
	require.Equal(t, "claude-sonnet-4-5-20250929", responseModel)

	headerMutation, bodyMutation, tokenUsage, responseModel, err = translator.ResponseBody(nil, strings.NewReader(responseTail), false, nil)
	require.NoError(t, err)
	require.Nil(t, headerMutation)
	require.Nil(t, bodyMutation)
	expected = tokenUsageFrom(10, 1, 16, 26)
	require.Equal(t, expected, tokenUsage)
	require.Equal(t, "claude-sonnet-4-5-20250929", responseModel)
}
